import logging
from typing import (
    Optional,
    Dict,
    List,
    Any
)
from datetime import datetime

class CitationManager:
    """改进的引用管理器"""
    
    def __init__(self) -> None:
        self.citations = []  # 引用列表
        self.citation_counter = 1  # 引用计数器
        self.paper_map = {}  # {paper_key: number}
        self.logger = logging.getLogger(__name__)
        self.topic = "default"  # 主题标识符
        self._version = 0  # 版本控制
        self._keys = set()  # 用于增量更新的键集合

    def add_citation(self, paper: dict) -> int:
        """添加单个引用，确保引用的唯一性和编号的连续性"""
        try:
            # 生成唯一键
            paper_key = self._generate_unique_key(paper)
            
            # 创建新引用
            citation = self._create_validated_citation(paper)
            if not citation:
                return 0
                
            # 检查是否已存在相同的引用（基于paper_key）
            if paper_key in self.paper_map:
                return self.paper_map[paper_key]
            
            # 分配新的引用编号
            next_number = self.citation_counter
            citation['number'] = next_number
            self.citations.append(citation)
            self.paper_map[paper_key] = next_number
            self.citation_counter += 1
            
            return citation['number']
        except Exception as e:
            self.logger.error(f"Citation addition failed: {str(e)}")
            return 0

    def update(self, new_papers):
        """增量更新机制"""
        version_changed = False
        for paper in new_papers:
            key = self._generate_unique_key(paper)
            if key not in self._keys:
                self.citations.append({
                    **paper,
                    "version": self._version
                })
                self._keys.add(key)
                version_changed = True
        if version_changed:
            self._version += 1  # 版本号递增触发更新
            
    def get_current_version(self):
        return self._version

    def _create_validated_citation(self, paper: dict) -> Optional[Dict[str, Any]]:
        """创建并验证引用"""
        try:
            # 补全必要字段
            required_fields = {
                'title': paper.get('title', f"Untitled_{len(self.citations) + 1}"),
                'venue': paper.get('venue', "Conf"),
                'year': paper.get('year', datetime.now().year),
                'chunk_id': paper.get('chunk_id', str(len(self.citations) + 1)),
                'authors': paper.get('authors', ['Anonymous']),
                'doi': paper.get('doi', ''),
                'topic': paper.get('topic', self.topic)  # 添加主题字段
            }
            
            # 验证并补全字段
            validated_paper = {}
            for field, default in required_fields.items():
                value = paper.get(field, default)
                if not value:
                    value = default
                validated_paper[field] = value

            return validated_paper
            
        except Exception as e:
            self.logger.error(f"引用创建失败: {str(e)}")
            return None

    def add_multiple_citations(self, paper: dict, count: int = 1) -> List[int]:
        """添加单个引用，为保持兼容性保留count参数但默认为1"""
        citation_nums = []
        # 只添加一个引用，忽略count参数
        if num := self.add_citation(paper):
            citation_nums.append(num)
        return citation_nums

    def _generate_unique_key(self, paper: dict) -> str:
        """生成论文唯一键，使用paper_id和chunk_id的组合"""
        try:
            paper_id = paper.get('paper_id', '')
            chunk_id = str(paper.get('chunk_id', ''))
            title = paper.get('title', '').strip()
            
            # 如果没有paper_id，使用标题作为替代
            if not paper_id:
                paper_id = title
                
            # 确保chunk_id非空
            if not chunk_id:
                chunk_id = 'default'
                
            # 使用paper_id和chunk_id的组合作为唯一键
            return f"{paper_id}_{chunk_id}"
        except Exception as e:
            self.logger.error(f"生成唯一键失败: {str(e)}")
            return f"error_key_{len(self.citations)}"

    def get_citations(self) -> List[Dict[str, Any]]:
        """返回完整的引用列表，包含所有必要的引用信息"""
        return [{
            'number': cite['number'],
            'title': cite['title'],
            'venue': cite['venue'],
            'year': cite['year'],
            'chunk_id': cite['chunk_id'],
            'authors': cite.get('authors', ['Anonymous']),
            'doi': cite.get('doi', ''),
            'abstract': cite.get('abstract', ''),
            'keywords': cite.get('keywords', []),
            'topic': cite.get('topic', self.topic)
        } for cite in sorted(self.citations, key=lambda x: x['number'])]

    def validate_citations(self) -> bool:
        """增强的引用验证，确保单编号引用格式"""
        try:
            if not self.citations:
                return False
                
            # 检查编号连续性
            numbers = sorted(cite['number'] for cite in self.citations)
            if numbers != list(range(1, len(numbers) + 1)):
                self.logger.warning("引用编号不连续")
                return False
                
            # 检查字段完整性
            for cite in self.citations:
                if not all(cite.get(field) for field in ['title', 'venue', 'year']):
                    self.logger.warning(f"引用数据不完整: {cite}")
                    return False
                    
            # 检查引用编号唯一性
            if len(numbers) != len(set(numbers)):
                self.logger.warning("存在重复的引用编号")
                return False
                    
            return True
            
        except Exception as e:
            self.logger.error(f"引用验证失败: {str(e)}")
            return False

    def reset(self):
        """重置管理器状态，保留主题信息"""
        current_topic = self.topic  # 保存当前主题
        self.citations = []
        self.citation_counter = 1
        self.paper_map = {}  # 重置paper_map
        self._keys = set()  # 重置键集合
        self._version = 0  # 重置版本号
        self.topic = current_topic  # 恢复主题信息

    def get_citation_by_chunk(self, chunk_id: str) -> Optional[Dict[str, Any]]:
        """通过chunk_id获取引用信息"""
        for cite in self.citations:
            if cite['chunk_id'] == chunk_id:
                return cite
        return None
        
    def set_topic(self, topic: str) -> None:
        """设置当前主题"""
        if topic and topic.strip():
            self.topic = topic.strip()
            self.logger.info(f"设置主题为: {self.topic}")
    def process_citations(self, content: str, papers: List[Dict]) -> tuple:
        """处理文本中的引用标记并返回处理后的内容和引用列表"""
        try:
            # 初始化引用列表
            citations = []
            # 为每篇论文添加引用
            for paper in papers:
                citation_number = self.add_citation(paper)
                if citation_number:
                    citations.append(paper)
            
            # 返回处理后的内容和引用列表
            return content, citations
            
        except Exception as e:
            self.logger.error(f"处理引用失败: {str(e)}")
            return content, []
    
